{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c1a795a4",
   "metadata": {},
   "source": [
    "# Launch ComfyUI with TensorRT on Brev!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbb2d691",
   "metadata": {},
   "source": [
    "Welcome!\n",
    "\n",
    "In this notebook, we'll walk through how to use ComfyUI with TensorRT. ComfyUI is a powerful GUI for Stable Diffusion models. It uses a node-based architecture that allows users to construct very complex image/video generation workflows by connecting different nodes, each representing a different function or operation!\n",
    "\n",
    "TensorRT engines contain state-of-the-art optimizations to performance efficiently on NVIDIA GPUs. Once the TensorRT engine is built, you can implement it in your ComfyUI workflow to take advantage of the ultra-fast performance on your RTX GPU!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16808297-1050-4866-90cf-6bd145360343",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Environment Setup\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "OPTIONS = {}\n",
    "\n",
    "WORKSPACE = 'ComfyUI'\n",
    "USE_COMFYUI_MANAGER = True  #@param {type:\"boolean\"}\n",
    "UPDATE_COMFY_UI = True  #@param {type:\"boolean\"}\n",
    "INSTALL_CUSTOM_NODES_DEPENDENCIES = True  #@param {type:\"boolean\"}\n",
    "OPTIONS['UPDATE_COMFY_UI'] = UPDATE_COMFY_UI\n",
    "OPTIONS['USE_COMFYUI_MANAGER'] = USE_COMFYUI_MANAGER\n",
    "OPTIONS['INSTALL_CUSTOM_NODES_DEPENDENCIES'] = INSTALL_CUSTOM_NODES_DEPENDENCIES"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ededb20-837d-447a-b968-489fca04802d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "if ! [ -d WORKSPACE ]; then\n",
    "    echo \"Starting initial setup of ComfyUI...\"\n",
    "    git clone https://github.com/comfyanonymous/ComfyUI\n",
    "    if [ -d ComfyUI ]; then\n",
    "        echo \"Repository cloned successfully into 'ComfyUI'.\"\n",
    "        cd ComfyUI\n",
    "    else\n",
    "        echo \"Failed to clone repository.\"\n",
    "    fi\n",
    "else\n",
    "    echo \"Workspace directory already exists.\"\n",
    "fi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77c2635a-a3fe-45e8-bc63-42d70661358d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if OPTIONS['UPDATE_COMFY_UI']:\n",
    "    !echo \"Updating ComfyUI...\"\n",
    "    !git pull\n",
    "\n",
    "!echo \"Installing dependencies...\"\n",
    "!pip install accelerate\n",
    "!pip install einops \n",
    "!pip install transformers\n",
    "!pip install safetensors>=0.3.0 \n",
    "!pip install aiohttp \n",
    "!pip install pyyaml \n",
    "!pip install Pillow \n",
    "!pip install scipy \n",
    "!pip install tqdm \n",
    "!pip install psutil\n",
    "!pip install torch==2.2.2 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121\n",
    "!pip install torchsde\n",
    "!pip install kornia>=0.7.1 spandrel\n",
    "!pip install typer\n",
    "!pip install diffusers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3c877a1",
   "metadata": {},
   "source": [
    "### Optional, run this if you want to download ComfyUI manager"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "478acf97-9201-4b6e-b9b8-83eb82894017",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "\n",
    "import os\n",
    "\n",
    "if not os.path.isfile(\"ComfyUI-Manager/check.sh\"):\n",
    "    !chmod 755 ComfyUI-Manager/check.sh\n",
    "if not os.path.isfile(\"ComfyUI-Manager/scan.sh\"):\n",
    "    !chmod 755 ComfyUI-Manager/scan.sh\n",
    "if not os.path.isfile(\"ComfyUI-Manager/node_db/dev/scan.sh\"):\n",
    "    !chmod 755 ComfyUI-Manager/node_db/dev/scan.sh\n",
    "if not os.path.isfile(\"ComfyUI-Manager/scripts/install-comfyui-venv-linux.sh\"):\n",
    "    !chmod 755 ComfyUI-Manager/scripts/install-comfyui-venv-linux.sh\n",
    "if not os.path.isfile(\"ComfyUI-Manager/scripts/install-comfyui-venv-win.bat\"):\n",
    "    !chmod 755 ComfyUI-Manager/scripts/install-comfyui-venv-win.bat\n",
    "\n",
    "if not os.path.isdir(\"ComfyUI-Manager\"):\n",
    "    print(\"-= Initial setup ComfyUI-Manager =-\")\n",
    "    !git clone https://github.com/ltdrdata/ComfyUI-Manager \n",
    "\n",
    "%cd ComfyUI-Manager\n",
    "!git pull\n",
    "\n",
    "%cd ..\n",
    "\n",
    "%cd ~/verb-workspace"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d41fe1d2-fc53-4c47-9ff6-0e744ae2de58",
   "metadata": {},
   "outputs": [],
   "source": [
    "import shutil\n",
    "import os\n",
    "\n",
    "# Move the ComfyUI-Manager folder to the custom_nodes directory\n",
    "source_manager_folder = '/root/verb-workspace/ComfyUI-Manager'\n",
    "# source_manager_folder = '/home/ubuntu/verb-workspace/ComfyUI-Manager' (if you run into error, try this)\n",
    "destination_manager_folder = '/root/verb-workspace/ComfyUI/custom_nodes/ComfyUI-Manager'\n",
    "# destination_manager_folder = '/home/ubuntu/verb-workspace/ComfyUI/custom_nodes/ComfyUI-Manager' (if you run into error, try this)\n",
    "\n",
    "if os.path.exists(destination_manager_folder):\n",
    "    shutil.rmtree(destination_manager_folder)\n",
    "\n",
    "shutil.move(source_manager_folder, destination_manager_folder)\n",
    "\n",
    "assert os.path.exists(destination_manager_folder), \"Folder move failed for ComfyUI-Manager\"\n",
    "print(f'Folder ComfyUI-Manager moved to {destination_manager_folder}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9b82b83-6c8c-4958-a85a-699ea60a2445",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "if 'INSTALL_CUSTOM_NODES_DEPENDENCIES' in OPTIONS and OPTIONS['INSTALL_CUSTOM_NODES_DEPENDENCIES']:\n",
    "    print(\"-= Install custom nodes dependencies =-\")\n",
    "    !pip install GitPython\n",
    "    !python /root/verb-workspace/ComfyUI/custom_nodes/ComfyUI-Manager/cm-cli.py restore-dependencies"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd3d62c3",
   "metadata": {},
   "source": [
    "### Time to download some models (safetensors + checkpoints)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b635f87a-223d-4cbc-838f-8c96bbc048c0",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!pip install huggingface_hub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99232352-eb07-48d7-ace1-f97a9df69a00",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!wget -c https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.ckpt -P ./ComfyUI/models/checkpoints/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2b01186-5526-43e0-9d09-63cc34d41c48",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!wget -c https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.safetensors -P ./ComfyUI/models/vae/"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a21e7115",
   "metadata": {},
   "source": [
    "#### We're going to build the TensorRT engine for the model below (SDXL Turbo)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b794ae1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffusers import StableDiffusionPipeline\n",
    "\n",
    "# Download the model\n",
    "!wget -O ./ComfyUI/models/checkpoints/sd_xl_turbo_1.0_fp16.safetensors https://huggingface.co/stabilityai/sdxl-turbo/resolve/main/sd_xl_turbo_1.0_fp16.safetensors\n",
    "\n",
    "# Load the model from the local path\n",
    "pipeline = StableDiffusionPipeline.from_single_file(\n",
    "    \"./ComfyUI/models/checkpoints/sd_xl_turbo_1.0_fp16.safetensors\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "521bf3db-79f1-4770-bbc1-a7cd7b28f08d",
   "metadata": {},
   "outputs": [],
   "source": [
    "!cd ./ComfyUI/custom_nodes && git clone https://github.com/comfyanonymous/ComfyUI_TensorRT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f904ecec-3056-47c4-8fd8-2da756be51c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "!cd ./ComfyUI/custom_nodes/ComfyUI_TensorRT && pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac988cc5",
   "metadata": {},
   "source": [
    "#### This next cell will pull out 2 .JSON files you'll need to:\n",
    "\n",
    "1) Build the TensorRT Engine for SDXL Turbo Model\n",
    "2) Use the engine in the ComfyUI Workflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "703ab48d",
   "metadata": {},
   "outputs": [],
   "source": [
    "source_folder_path = '/root/verb-workspace/ComfyUI/custom_nodes/ComfyUI_TensorRT/workflows'\n",
    "# source_folder_path = '/home/ubuntu/verb-workspace/ComfyUI/custom_nodes/ComfyUI_TensorRT/workflows' (if you run into error, try this)\n",
    "destination_folder_path = '/root/verb-workspace/'\n",
    "# destination_folder_path = '/home/ubuntu/verb-workspace' (if you run into error, try this)\n",
    "\n",
    "json_files = [\n",
    "    'Create_SDXL_Turbo_TRT_Static.json',\n",
    "    'Build.TRT.Engine_SDXL_Turbo_Static.json'\n",
    "]\n",
    "\n",
    "# Move the workflows folder to the main directory\n",
    "destination_workflows_path = os.path.join(destination_folder_path, 'workflows')\n",
    "if os.path.exists(destination_workflows_path):\n",
    "    shutil.rmtree(destination_workflows_path)\n",
    "\n",
    "shutil.move(source_folder_path, destination_folder_path)\n",
    "\n",
    "# Move the JSON files out of the workflows folder into the main directory\n",
    "for json_file in json_files:\n",
    "    source_json_path = os.path.join(destination_workflows_path, json_file)\n",
    "    destination_json_path = os.path.join(destination_folder_path, json_file)\n",
    "    shutil.move(source_json_path, destination_json_path)\n",
    "    \n",
    "    assert os.path.exists(destination_json_path), f\"File move failed for {json_file}\"\n",
    "    print(f'File {json_file} moved to {destination_json_path}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f2b03f4-90fe-467f-b020-1370ab819274",
   "metadata": {},
   "outputs": [],
   "source": [
    "import subprocess\n",
    "import threading\n",
    "import time\n",
    "import socket\n",
    "\n",
    "def check_server(port):\n",
    "    while True:\n",
    "        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:\n",
    "            result = sock.connect_ex(('127.0.0.1', port))\n",
    "            if result == 0:\n",
    "                print(\"\\nServer is up and running at http://129.80.69.99:{}\".format(port))\n",
    "                break\n",
    "            time.sleep(0.5)\n",
    "\n",
    "threading.Thread(target=check_server, daemon=True, args=(8188,)).start()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35d97a31-1697-4ab3-82fb-d81fa7d43963",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!python /root/verb-workspace/ComfyUI/main.py --listen 0.0.0.0 --port 8188 --dont-print-server --disable-xformers\n",
    "# !python /home/ubuntu/verb-workspace/ComfyUI/main.py --listen 0.0.0.0 --port 8188 --dont-print-server --disable-xformers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c302da0",
   "metadata": {},
   "source": [
    "### We just started the ComfyUI server! \n",
    "\n",
    "##### We can access it by going to the Brev instance page and clicking on the Port 8188 URL\n",
    "\n",
    "![image.png](https://uohmivykqgnnbiouffke.supabase.co/storage/v1/object/public/notebooks/Screenshot_2024-06-04_at_13.00.59.png?t=2024-06-04T20%3A01%3A47.076Z)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afbb2db6",
   "metadata": {},
   "source": [
    "#### Download the 2 JSON files we pulled into the main directory to your local computer \n",
    "##### Files you need to download: 'Create_SDXL_Turbo_TRT_Static.json', 'Build.TRT.Engine_SDXL_Turbo_Static.json'\n",
    "\n",
    "##### The ComfyUI server should be running, so to begin, click on the URL at Port 8188 (as shown above) \n",
    "\n",
    "##### Once you see the default workflow, we can use the `Load` Button to bring in the 'Build.TRT.Engine_SDXL_Turbo_Static.json' file from your laptop. \n",
    "\n",
    "![image.png](https://uohmivykqgnnbiouffke.supabase.co/storage/v1/object/public/notebooks/LoadBuildjson.png?t=2024-06-04T18%3A01%3A20.036Z)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4691c25",
   "metadata": {},
   "source": [
    "##### Make sure to upload the 'Build.TRT.Engine_SDXL_Turbo_Static.json' file from your laptop now!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd213a93",
   "metadata": {},
   "source": [
    "##### Click `Queue Prompt` to begin the workflow. This will build the TensorRT engine (takes a couple minutes)\n",
    "\n",
    "![image.png](https://uohmivykqgnnbiouffke.supabase.co/storage/v1/object/public/notebooks/BuildTRTEngine.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7725c80f",
   "metadata": {},
   "source": [
    "##### Now that the engine is built, it's time to actually run a ComfyUI workflow with it and see how fast it is yourself. Click the `Load` Button again to bring in the 'Create_SDXL_Turbo_TRT_Static.json' file. \n",
    "\n",
    "![image.png](https://uohmivykqgnnbiouffke.supabase.co/storage/v1/object/public/notebooks/loadCreatejson.png?t=2024-06-04T18%3A13%3A54.983Z)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f243ea74",
   "metadata": {},
   "source": [
    "\n",
    "##### Now, click `Queue Prompt` to run this workflow and generate images with SDXL Turbo!\n",
    "\n",
    "![image.png](https://uohmivykqgnnbiouffke.supabase.co/storage/v1/object/public/notebooks/TRTComfyUIWorkflow.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9f58b03",
   "metadata": {},
   "source": [
    "##### Here was my output using the TRT Engine in my ComfyUI workflow!\n",
    "\n",
    "![image.png](https://uohmivykqgnnbiouffke.supabase.co/storage/v1/object/public/notebooks/Screenshot_2024-06-04_at_11.15.57.png?t=2024-06-04T18%3A16%3A50.463Z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1b6853f-af32-4543-80aa-9b86b275d42a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def check_server(port):\n",
    "    while True:\n",
    "        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:\n",
    "            result = sock.connect_ex(('127.0.0.1', port))\n",
    "            if result == 0:\n",
    "                print(\"\\nServer is up and running at http://129.80.69.99:{}\".format(port))\n",
    "                break\n",
    "            time.sleep(0.5)\n",
    "\n",
    "threading.Thread(target=check_server, daemon=True, args=(8188,)).start()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
